import os

from flask import Flask, request, jsonify
from flask_cors import CORS
from transformers import AutoModelWithLMHead, AutoTokenizer, pipeline

app = Flask(__name__)
CORS(app)
app.BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
import warnings

warnings.filterwarnings('ignore')
# 所有模型
models = {}


def load_model():
    print("正在加载 [zh-en] 模型")
    zh_en_model = AutoModelWithLMHead.from_pretrained("Helsinki-NLP/opus-mt-zh-en")
    zh_en_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-zh-en")
    zh_en_translation = pipeline('translation_zh_to_en', model=zh_en_model, tokenizer=zh_en_tokenizer)
    models['zh-en'] = zh_en_translation

    print("正在加载 [en-zh] 模型")

    en_zh_model = AutoModelWithLMHead.from_pretrained("Helsinki-NLP/opus-mt-en-zh", mirror='tuna')
    en_zh_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-zh", mirror='tuna')
    en_zh_translation = pipeline('translation_en_to_zh', model=en_zh_model, tokenizer=en_zh_tokenizer)
    models['en-zh'] = en_zh_translation

    print("加载完成")
    return models


@app.route("/translation/<path:trans_type>", methods=['POST', 'GET'])
def translation(trans_type):
    translation_model = models[trans_type]

    content = (request.json if request.method == 'POST' else request.args)['content']

    res = translation_model(content)[0]['translation_text']
    return jsonify({
        "data": res,
    })


if __name__ == '__main__':
    print("加载模型中")
    load_model()
    app.run('0.0.0.0', 2222)
